'''
LlcFilters

Filters used when dealing with OpenMP pragmas

Created on 26/10/2010

:author: rreyes
'''

from yacf.Tools.Tree import NodeNotFound

from yacf.Backends.Common.Visitors.GenericVisitors import GenericFilterVisitor,\
    AbstractReverseVisitor
from yacf.Tools.Declarations import decl_of_id
from yacf.Backends.C.Exceptions import NonOmpCompilant
from yacf.Frontend.C99 import c99_ast as c_ast 
from yacf.Frontend.Acc import acc_ast 

from yacf.Frontend.InternalRepr import AstToIR
from yacf.Backends.C99.Writers.OmpWriter import OmpWriter
from yacf.Frontend.SymbolTable import SymbolTable, SymbolTableBuilder, getSymbolTableFromNode, IdentifierNotFound

import copy

class AccKernelsReverse(AbstractReverseVisitor):
    def __init__(self, ):
        self._clauses = {}
        self._rootNode = None
        super(AccKernelsReverse, self).__init__(condition_func = lambda x : isinstance(x, acc_ast.AccKernels))

class AccAnyRegionReverse(AbstractReverseVisitor):
    def __init__(self, ):
        self._clauses = {}
        self._rootNode = None

        def condition(node):
            if isinstance(node, acc_ast.AccKernels) \
                    or isinstance(node, acc_ast.AccParallel) \
                    or isinstance(node, acc_ast.AccData):
                return True
            return False

        super(AccAnyRegionReverse, self).__init__(condition_func = condition)


class FuncDefReverse(AbstractReverseVisitor):
    def __init__(self, ):
        self._clauses = {}
        self._rootNode = None
        super(FuncDefReverse, self).__init__(condition_func = lambda x : isinstance(x, c_ast.FuncDef))




class AccLoopFilter(GenericFilterVisitor):
    """ Returns a LlcFor node, the parallel container and the function container

         By defining specific visitor methods for *FuncDef* and *LlcParallel*, we can
            save the last node visited of this types. Giving the fact that the visit is
            done in syntax order, the last visited node will be the previous (parent) node
            of the wanted node.

    """
    
    def __init__(self, prev_brother = None):
        self._parentRegion = None
        self._funcdef = None
        def condition(node):
            """ LlcFor filter """
            return type(node) == acc_ast.AccLoop
        super(AccLoopFilter, self).__init__(condition_func = condition, 
                prev_brother = prev_brother)

    def visit_FuncDef(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._funcdef = node
        return self.generic_visit(node, offset, ignore)

    def visit_AccKernels(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._parentRegion = node
        return self.generic_visit(node, offset, ignore)

    def visit_AccParallel(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._parentRegion = node
        return self.generic_visit(node, offset, ignore)


  
class AccAbstractRegionFilter(GenericFilterVisitor):

    def visit_FuncDef(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._funcdef = node
        return self.generic_visit(node, offset, ignore)

    def visit_AccKernels(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._parentRegion = node
        return self.generic_visit(node, offset, ignore)

    def visit_AccParallel(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._parentRegion = node
        return self.generic_visit(node, offset, ignore)

    def visit_AccData(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._parentRegion = node
        return self.generic_visit(node, offset, ignore)

class AccHostDataFilter(AccAbstractRegionFilter):
    """ 

    """
    def __init__(self, prev_brother = None):
        self._parentRegion = None
        self._funcdef = None
        def condition(node):
            """ AccData filter """
            return type(node) == acc_ast.AccHostData
        super(AccHostDataFilter, self).__init__(condition_func = condition, 
                prev_brother = prev_brother)
 


class AccUpdateFilter(AccAbstractRegionFilter):
    """ 

    """
    def __init__(self, prev_brother = None):
        self._parentRegion = None
        self._funcdef = None
        def condition(node):
            """ AccUpdate filter """
            return isinstance(node, acc_ast.AccUpdate)
        super(AccUpdateFilter, self).__init__(condition_func = condition, 
                prev_brother = prev_brother)


class AccWaitFilter(AccAbstractRegionFilter):
    """ 

    """
    def __init__(self, prev_brother = None):
        self._parentRegion = None
        self._funcdef = None
        def condition(node):
            """ AccUpdate filter """
            return isinstance(node, acc_ast.AccWait)
        super(AccWaitFilter, self).__init__(condition_func = condition, 
                prev_brother = prev_brother)




class AccKernelsFilter(GenericFilterVisitor):
    """ Returns a LlcRegion node, the parallel container and the function 
        container

         By defining specific visitor methods for *FuncDef* and *LlcRegion*, we can
            save the last node visited of this types. Giving the fact that the visit is
            done in syntax order, the last visited node will be the previous (parent) 
            node of the wanted node.


    """
 
    def __init__(self, prev_brother = None,):
        self._parentRegion = None
        self._funcdef = None
        def condition_func(node):
            if isinstance(node, acc_ast.AccKernels):
                return True
            return False
        super(AccKernelsFilter, self).__init__(condition_func = condition_func,
                                                    prev_brother = prev_brother,)


    def visit_FuncDef(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._funcdef = node
        return self.generic_visit(node, offset, ignore)
    
    def get_func_def(self):
        return self._funcdef




class AccDataFilter(GenericFilterVisitor):
    """ Returns a LlcRegion node, the parallel container and the function 
        container

         By defining specific visitor methods for *FuncDef* and *LlcRegion*, we can
            save the last node visited of this types. Giving the fact that the visit is
            done in syntax order, the last visited node will be the previous (parent) 
            node of the wanted node.


    """
 
    def __init__(self, prev_brother = None,):
        self._parentRegion = None
        self._funcdef = None
        def condition_func(node):
            if isinstance(node, acc_ast.AccData):
                return True
            return False
        super(AccDataFilter, self).__init__(condition_func = condition_func,
                                                    prev_brother = prev_brother,)


    def visit_FuncDef(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._funcdef = node
        return self.generic_visit(node, offset, ignore)
    
    def get_func_def(self):
        return self._funcdef


class AccParallelFilter(GenericFilterVisitor):
    """ Returns a AccParalle node
    """
 
    def __init__(self, prev_brother = None,):
        self._parentRegion = None
        self._funcdef = None
        def condition_func(node):
            if isinstance(node, acc_ast.AccParallel):
                return True
            return False
        super(AccParallelFilter, self).__init__(condition_func = condition_func,
                                                    prev_brother = prev_brother,)



class AccRegionsFilter(GenericFilterVisitor):
    """ Returns a LlcRegion node, the parallel container and the function 
        container

         By defining specific visitor methods for *FuncDef* and *LlcRegion*, we can
            save the last node visited of this types. Giving the fact that the visit is
            done in syntax order, the last visited node will be the previous (parent) 
            node of the wanted node.


    """
 
    def __init__(self, prev_brother = None,):
        self._parentRegion = None
        self._funcdef = None
        def condition_func(node):
            if isinstance(node, acc_ast.AccKernels) or \
                    isinstance(node, acc_ast.AccParallel):
                return True
            return False
        super(AccRegionsFilter, self).__init__(condition_func = condition_func,
                                                    prev_brother = prev_brother,)


    def visit_FuncDef(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._funcdef = node
        return self.generic_visit(node, offset, ignore)
    
    def get_func_def(self):
        return self._funcdef



class AccScopeReverse(AbstractReverseVisitor):
    def __init__(self, rootNode = None):
        self._rootNode = rootNode
        # It is safe to make the following assumptions at this point, as AccScopeReverse 
        #   is only called from an IR-1
        # 1.- All nodes have been linked/relinked
        # 2. 
        if rootNode:
           self._symbolTable = getSymbolTableFromNode(rootNode)
        else:
           # Code to build the TS should be here
           raise NotImplementedError
        self._clauses = {}
        super(AccScopeReverse, self).__init__(condition_func = lambda x : isinstance(x, c_ast.FileAST))

    def _get_dict_from_clauses(self, clauses,):
        """ Return a dict of clauses from a list of LlcClause objects
            
              Example: [LlcClause('REDUCTION', ...), LlcClause('PRIVATE', ...)]
                 will return:  {'REDUCTION' : [....] , 'PRIVATE' : [...]}

            :return: dict with clauses
        """
        fileAst = self._rootNode
        mem_clauses = ['COPYIN', 'COPYOUT', 'COPY', 'CREATE', 'PCOPYIN','PCOPYOUT','PCOPY','PCREATE',  'USEDEVICE', 'PRIVATE',]
        clause_names = ['NOWAIT', 'REDUCTION',  'COLLAPSE', 'NAME','ASYNC'] + mem_clauses
        clause_dict = {}
        
        # Note: Each identifiers is a ParamList
        for elem in clauses:
            if not clause_dict.has_key(elem.name):
                    clause_dict[elem.name] = []
            # Clauses with declarations
            if elem.name == 'REDUCTION':
                id_list = []
                for id in elem.identifiers.params:
                    try:
                        symbol = self._symbolTable.lookUp(id)
                    # If a declaration with the same name is already stored, pass. 
                    #  Otherwise, append it to the list
                        for stored_symbol in clause_dict[elem.name]:
                            if symbol.name == stored_symbol.name:
                                break
                        else:
                            id_list.append(symbol)
                    except IdentifierNotFound:
                            raise NonOmpCompilant(" Declaration of " + id.name + " in " + elem.name + " clause could not be found ")

                # Add REDUCTION to the dictionary
                clause_dict[elem.name].append((str(elem.type),id_list, ))
            elif elem.name in mem_clauses:
                for id in elem.identifiers.params:
                    try:
                            symbol = self._symbolTable.lookUp(id)
                            self._symbolTable.lookUp(id).extra['startAddr'] = None
                            if type(elem.identifiers) == acc_ast.AccParamList \
                                and elem.identifiers.definitions.has_key(id):
                               current = elem.identifiers.definitions[id]
                               # This element has been redefined at clause, 
                               #  we have to alter its symbol adding the extra information
                               sizeof_ast = c_ast.UnaryOp(op = 'sizeof', expr = None)   
                               sizeof_ast.expr = c_ast.Constant(value = symbol.btype, type = 'string')
                               if type(current) == type(tuple()):
                                   print " CURRENT " + str(id) + " is a 2D array "
                                   sizeExpr_d1_subtree = c_ast.BinaryOp(op = '*', 
                                                            left = current[1].start, 
                                                            right = sizeof_ast)
                                   sizeExpr_d2_subtree = c_ast.BinaryOp(op = '*', 
                                                            left = current[0].start, 
                                                            right = sizeExpr_d1_subtree) 
                                   self._symbolTable.lookUp(id).sizeExpression = sizeExpr_d2_subtree
                                   # TODO: ENABLE/DISABLE THIS FEATURE
                                   from yacf.Frontend.SymbolTable import MatrixInfo
                                   mInfo = None
                                   if current[1].stride:
                                       print " Stride "
                                       mInfo = MatrixInfo(current[0].start, current[1].start, 
                                                          sizeof_ast.expr, current[1].stride)
                                   else:
                                       # Only for non-contiguous arrays
                                       mInfo = MatrixInfo(current[0].start, current[1].start, 
                                                          sizeof_ast.expr, current[1].start)
            
                                   self._symbolTable.lookUp(id).extra['matrixInfo'] =  mInfo
                                   if current[0].end:
                                        self._symbolTable.lookUp(id).extra['startAddr'] = current[0].start
                               else:
                                   if current.end and current.start:
                                       sizeExpr_subtree = c_ast.BinaryOp(op = '*', 
                                                        left = c_ast.BinaryOp(op='-',
                                                                left = current.end,
                                                                right = current.start,), 
                                                        right = sizeof_ast)
                                   else:
                                       sizeExpr_subtree =  c_ast.BinaryOp(op = '*', 
                                                                left = current.start, 
                                                                right = sizeof_ast)
                                   self._symbolTable.lookUp(id).sizeExpression = sizeExpr_subtree
                                   self._symbolTable.lookUp(id).extra['matrixInfo'] =  None
                                   if current.end:
                                        self._symbolTable.lookUp(id).extra['startAddr'] =  current.start
                    # If a declaration with the same name is already stored, pass. Otherwise, append it to the list
                            for stored_symbol in clause_dict[elem.name]:
                                if symbol.name == stored_symbol.name:
                                    break
                            else:
                                clause_dict[elem.name].append(symbol)
                    except IdentifierNotFound:
                        print " Variable declaration for "  + id.name + " in " + elem.name + " clause in current scope "
                        tutu = self._symbolTable.lookUp(id)
                        raise NonOmpCompilant(" Variable declaration for "  + id.name + " in " + elem.name + " clause in current scope ")

            elif elem.name == 'NOWAIT':
                clause_dict[elem.name] = True
            elif elem.name == 'COLLAPSE':
                clause_dict[elem.name] = int(elem.identifiers.value)
            elif elem.name == 'NAME':
                clause_dict[elem.name] = elem.label
            elif elem.name == 'ASYNC':
                clause_dict[elem.name] = elem.value


        # Add empty keywords
        for name in clause_names:
            if not clause_dict.has_key(name):
                clause_dict[name] = []
        return  clause_dict
    
    
    def _merge_dict_per_key(self, destination, source):
        """ Merge two clause dicts per key value. 
            Sample:
            .. ipython:
            
            [0] a = { 'A' : [0,1], 'B' : [2,], 'C' : []}
            [1] b = { 'B' : [3,] , 'C' : 3 }
            [2] c = _merge_dict_per_key(a, b)
            [3] c = {'A' : [0,1], 'B' : [2,3], 'C' : 3}
            
        """
        source_cpy = copy.copy(source)
        for elem in destination:
            if source_cpy.has_key(elem):
                if type(destination[elem]) == type([]) and (len(destination[elem]) > 0):
                    destination[elem] += source[elem]
                else:
                    if type(source[elem]) == type([]) and (len(source[elem]) > 0):
                        destination[elem] = source[elem]
                    elif type(source[elem]) != type([]):
                        destination[elem] = source[elem]
                    else:
                        # Keep destination
                        pass
                del source_cpy[elem]
            
        destination.update(source_cpy)
        
        return destination        

    
    def visit_Compound(self, node):
        """ Compound statements may contain other target clauses, look for them """
#        for elem in node.block_items:
#            if isinstance(elem, c_ast.Pragma):
#                if isinstance(elem.stmt, omp_ast.LlcTargetDevice):
#                    self._clauses = self._merge_dict_per_key(self._clauses, self._get_dict_from_clauses(elem.stmt.clauses))
        return self.generic_visit(node)
        

    def visit_AccKernels(self, node):
        #=======================================================================
        # print "ACC KERNELS "  + str(self._clauses) + " Node: " + str([str(elem) for elem in node.clauses])
        #=======================================================================
        # self._clauses.update(self._get_dict_from_clauses(node.clauses))
        self._merge_dict_per_key(self._clauses, self._get_dict_from_clauses(node.clauses))
        return self.generic_visit(node)

    def visit_AccLoop(self, node):
        #=======================================================================
        # print "ACC LOOP " + str(self._clauses) + " Node: " + str([str(elem) for elem in node.clauses])
        #=======================================================================
        # self._clauses.update(self._get_dict_from_clauses(node.clauses))
        self._merge_dict_per_key(self._clauses, self._get_dict_from_clauses(node.clauses))
        return self.generic_visit(node)


    def visit_AccData(self, node):
        #=======================================================================
        # print "ACC DATA " + str(self._clauses) + " Node: " + str([str(elem) for elem in node.clauses])
        #=======================================================================
        # self._clauses.update(self._get_dict_from_clauses(node.clauses))
        self._merge_dict_per_key(self._clauses, self._get_dict_from_clauses(node.clauses))
        return self.generic_visit(node)


    def get_clauses(self):
        return self._clauses




